import math
import os
import torch.nn as nn
import torch.nn.functional as F
import torch
import matplotlib.pyplot as plt
from tools.data_generator import generate_logic_data
from tools.remove_temp_file import remove_path
from bdtime import tt
from tools.my_trainer_and_test import train, test, MODEL_SAVE_DIR


class MultiLayerLogisticClassification(nn.Module):

    def __init__(self, in_features, out_features, middle_features=10, bias=True):
        super().__init__()

        self.fc_1 = nn.Linear(in_features, middle_features, bias=bias)
        self.fc_2 = nn.Linear(middle_features, out_features, bias=bias)

        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.relu(self.fc_1(x))
        output = self.sigmoid(self.fc_2(x))
        return output


class LogisticClassification(nn.Module):
    def __init__(self, in_features, out_features, bias=1):
        super().__init__()
        self.fc = nn.Linear(in_features, out_features, bias=bias)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        output = self.sigmoid(self.fc(x))
        return output


class LogisticRegression(nn.Module):
    def __init__(self, in_features, out_features, bias=1):
        super().__init__()
        self.fc = nn.Linear(in_features, out_features, bias=bias)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        output = self.sigmoid(self.fc(x))
        return output


if __name__ == '__main__':
    from tools import data_generator
    from tools.change_cwd_to_main import change_cwd_to_main
    change_cwd_to_main()

    noise_config = None
    # noise_config = (0, 1)
    x_data, y_data = data_generator.generate_logic_data(size=10000, noise_config=noise_config)
    x_train, x_test, y_train, y_test = data_generator.split_train_and_test(x_data, y_data, test_size=0.2)

    bias = True

    # model = LogisticRegression(in_features=2, out_features=1, bias=bias)
    # experiment_name = 'generate_logic_data__LogisticRegression'

    model = MultiLayerLogisticClassification(in_features=2, out_features=1, middle_features=10, bias=bias)
    experiment_name = 'generate_logic_data__MultiLayerLogisticClassification'

    # --- loss_function and optimizer
    criterion = nn.BCELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

    # x_train, x_test, y_train, y_test
    train((x_train, x_test, y_train, y_test), model, optimizer, criterion,
          total_epoch=20000, log_interval=0.1,
          experiment_name=experiment_name, save=True)
